Populating next right pointers in each node

Time: O(N); Space: O(1); medium

You are given a perfect binary tree where all leaves are on the same level, and every parent has two children.

The binary tree has the following definition:

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.next = None

Populate each next pointer to point to its next right node.

If there is no next right node, the next pointer should be set to NULL.

Initially, all next pointers are set to NULL.

Follow up:

• You may only use constant extra space.

• Recursive approach is fine, you may assume implicit stack space does not count as extra space for this problem.

Example 1:

Input: root = TreeNode [1,2,3,4,5,6,7]

Output: [1,#,2,3,#,4,5,6,7,#]

Explanation:

• Given the above perfect binary tree (Figure A), your function should populate each next pointer to point to its next right node, just like in Figure B.

• The serialized output is in level order as connected by the next pointers, with ‘#’ signifying the end of each level.

Constraints:

• The number of nodes in the given tree is less than 4096.

• -1000 <= node.val <= 1000

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.next = None

    def __repr__(self):
        if self is None:
            return "None"
        else:
            return "{} -> {}".format(self.val, repr(self.next))

class Solution1(object):
    """
    Time: O(N)
    Space: O(1)
    """
    def connect(self, root):
        """
        :type root: TreeNode
        :rtype: nothing
        """
        head = root
        while head:
            cur = head
            while cur and cur.left:
                cur.left.next = cur.right
                if cur.next:
                    cur.right.next = cur.next.left
                cur = cur.next
            head = head.left

s = Solution1()
root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.left = TreeNode(6)
root.right.right = TreeNode(7)
s.connect(root)

print(root)              # 1 -> None
print(root.left)         # 2 -> 3 -> None
print(root.right)        # 3 -> None
print(root.left.left)    # 4 -> 5 -> 6 -> 7 -> None
print(root.left.right)   # 5 -> 6 -> 7 -> None
print(root.right.left)   # 6 -> 7 -> None
print(root.right.right)  # 7 -> None

1 -> None
2 -> 3 -> None
3 -> None
4 -> 5 -> 6 -> 7 -> None
5 -> 6 -> 7 -> None
6 -> 7 -> None
7 -> None

2. Recusion

class Solution2(object):
    def connect(self, root):
        """
        :type root: TreeNode
        :rtype: nothing
        """
        if root is None:
            return

        if root.left:
            root.left.next = root.right

        if root.right and root.next:
            root.right.next = root.next.left

        self.connect(root.left)
        self.connect(root.right)

s = Solution2()
root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)
root.right.left = TreeNode(6)
root.right.right = TreeNode(7)
s.connect(root)

print(root)              # 1 -> None
print(root.left)         # 2 -> 3 -> None
print(root.right)        # 3 -> None
print(root.left.left)    # 4 -> 5 -> 6 -> 7 -> None
print(root.left.right)   # 5 -> 6 -> 7 -> None
print(root.right.left)   # 6 -> 7 -> None
print(root.right.right)  # 7 -> None

1 -> None
2 -> 3 -> None
3 -> None
4 -> 5 -> 6 -> 7 -> None
5 -> 6 -> 7 -> None
6 -> 7 -> None
7 -> None

See also:

https://leetcode.com/problems/populating-next-right-pointers-in-each-node

https://www.lintcode.com/problem/populating-next-right-pointers-in-each-node/description